{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import re\n",
    "from IPython.core.debugger import set_trace\n",
    "from pprint import pprint\n",
    "import unicodedata\n",
    "from transformers import BertModel, BertTokenizerFast\n",
    "import copy\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import torch.optim as optim\n",
    "import glob\n",
    "import time\n",
    "import logging\n",
    "from common.utils import Preprocessor, DefaultLogger\n",
    "from tplinker_plus import (HandshakingTaggingScheme,\n",
    "                          DataMaker4Bert, \n",
    "                          DataMaker4BiLSTM, \n",
    "                          TPLinkerPlusBert, \n",
    "                          TPLinkerPlusBiLSTM,\n",
    "                          MetricsCalculator)\n",
    "import wandb\n",
    "from glove import Glove\n",
    "import numpy as np\n",
    "import config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = config.train_config\n",
    "hyper_parameters = config[\"hyper_parameters\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(config[\"device_num\"])\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for reproductivity\n",
    "torch.manual_seed(hyper_parameters[\"seed\"]) # pytorch random seed\n",
    "torch.backends.cudnn.deterministic = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_home = config[\"data_home\"]\n",
    "experiment_name = config[\"exp_name\"]    \n",
    "train_data_path = os.path.join(data_home, experiment_name, config[\"train_data\"])\n",
    "valid_data_path = os.path.join(data_home, experiment_name, config[\"valid_data\"])\n",
    "rel2id_path = os.path.join(data_home, experiment_name, config[\"rel2id\"])\n",
    "ent2id_path = os.path.join(data_home, experiment_name, config[\"ent2id\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if config[\"logger\"] == \"wandb\":\n",
    "    # init wandb\n",
    "    wandb.init(project = experiment_name, \n",
    "               name = config[\"run_name\"],\n",
    "               config = hyper_parameters # Initialize config\n",
    "              )\n",
    "\n",
    "    wandb.config.note = config[\"note\"]          \n",
    "\n",
    "    model_state_dict_dir = wandb.run.dir\n",
    "    logger = wandb\n",
    "else:\n",
    "    logger = DefaultLogger(config[\"log_path\"], experiment_name, config[\"run_name\"], config[\"run_id\"], hyper_parameters)\n",
    "    model_state_dict_dir = config[\"path_to_save_model\"]\n",
    "    if not os.path.exists(model_state_dict_dir):\n",
    "        os.makedirs(model_state_dict_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = json.load(open(train_data_path, \"r\", encoding = \"utf-8\"))\n",
    "valid_data = json.load(open(valid_data_path, \"r\", encoding = \"utf-8\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @specific\n",
    "if config[\"encoder\"] == \"BERT\":\n",
    "    tokenizer = BertTokenizerFast.from_pretrained(config[\"bert_path\"], add_special_tokens = False, do_lower_case = False)\n",
    "    tokenize = tokenizer.tokenize\n",
    "    get_tok2char_span_map = lambda text: tokenizer.encode_plus(text, return_offsets_mapping = True, add_special_tokens = False)[\"offset_mapping\"]\n",
    "elif config[\"encoder\"] in {\"BiLSTM\", }:\n",
    "    tokenize = lambda text: text.split(\" \")\n",
    "    def get_tok2char_span_map(text):\n",
    "        tokens = text.split(\" \")\n",
    "        tok2char_span = []\n",
    "        char_num = 0\n",
    "        for tok in tokens:\n",
    "            tok2char_span.append((char_num, char_num + len(tok)))\n",
    "            char_num += len(tok) + 1 # +1: whitespace\n",
    "        return tok2char_span"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preprocessor = Preprocessor(tokenize_func = tokenize, \n",
    "                            get_tok2char_span_map_func = get_tok2char_span_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train and valid max token num\n",
    "max_tok_num = 0\n",
    "all_data = train_data + valid_data \n",
    "    \n",
    "for sample in all_data:\n",
    "    tokens = tokenize(sample[\"text\"])\n",
    "    max_tok_num = max(max_tok_num, len(tokens))\n",
    "max_tok_num"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if max_tok_num > hyper_parameters[\"max_seq_len\"]:\n",
    "    train_data = preprocessor.split_into_short_samples(train_data, \n",
    "                                                          hyper_parameters[\"max_seq_len\"], \n",
    "                                                          sliding_len = hyper_parameters[\"sliding_len\"], \n",
    "                                                          encoder = config[\"encoder\"]\n",
    "                                                         )\n",
    "    valid_data = preprocessor.split_into_short_samples(valid_data, \n",
    "                                                          hyper_parameters[\"max_seq_len\"], \n",
    "                                                          sliding_len = hyper_parameters[\"sliding_len\"], \n",
    "                                                          encoder = config[\"encoder\"]\n",
    "                                                         )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"train: {}\".format(len(train_data)), \"valid: {}\".format(len(valid_data)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# count_neg = 0 # 74.8% are neg samples 0.7485367594575303\n",
    "# for example in train_data + valid_data:\n",
    "#     if len(example[\"relation_list\"]) == 0 and len(example[\"entity_list\"]) == 0:\n",
    "#         count_neg += 1\n",
    "# print(count_neg/len(indexed_train_data + indexed_valid_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tagger (Decoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_seq_len = min(max_tok_num, hyper_parameters[\"max_seq_len\"])\n",
    "rel2id = json.load(open(rel2id_path, \"r\", encoding = \"utf-8\"))\n",
    "ent2id = json.load(open(ent2id_path, \"r\", encoding = \"utf-8\"))\n",
    "handshaking_tagger = HandshakingTaggingScheme(rel2id, max_seq_len, ent2id)\n",
    "tag_size = handshaking_tagger.get_tag_size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_equal_to(sample1, sample2):\n",
    "    assert sample1[\"id\"] == sample2[\"id\"]\n",
    "    assert sample1[\"text\"] == sample2[\"text\"]\n",
    "    memory_set = set()\n",
    "    for rel in sample2[\"relation_list\"]:\n",
    "        memory = \"{}\\u2E80{}\\u2E80{}\\u2E80{}\\u2E80{}\\u2E80{}\\u2E80{}\".format(rel[\"subject\"], \n",
    "                                                             rel[\"predicate\"], \n",
    "                                                             rel[\"object\"], \n",
    "                                                             *rel[\"subj_tok_span\"], \n",
    "                                                             *rel[\"obj_tok_span\"])\n",
    "        memory_set.add(memory)\n",
    "    for rel in sample1[\"relation_list\"]:\n",
    "        memory = \"{}\\u2E80{}\\u2E80{}\\u2E80{}\\u2E80{}\\u2E80{}\\u2E80{}\".format(rel[\"subject\"], \n",
    "                                                             rel[\"predicate\"], \n",
    "                                                             rel[\"object\"], \n",
    "                                                             *rel[\"subj_tok_span\"], \n",
    "                                                             *rel[\"obj_tok_span\"])\n",
    "        if memory not in memory_set:\n",
    "            set_trace()\n",
    "            return False\n",
    "    return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # check tagging and decoding\n",
    "# batch_size = hyper_parameters[\"batch_size\"]\n",
    "# for idx in tqdm(range(0, len(train_data), batch_size), desc = \"check tagging and decoding\"):\n",
    "#     batch_matrix_spots = []\n",
    "#     batch_data = train_data[idx:idx + batch_size]\n",
    "#     for sample in batch_data:\n",
    "#         matrix_spots = handshaking_tagger.get_spots(sample)\n",
    "# #         %timeit shaking_tagger.get_spots(sample)\n",
    "#         batch_matrix_spots.append(matrix_spots)\n",
    "    \n",
    "#     # tagging\n",
    "#     # batch_shaking_tag: (batch_size, rel_id, seq_len, seq_len)\n",
    "#     batch_shaking_tag = handshaking_tagger.spots2shaking_tag4batch(batch_matrix_spots)\n",
    "# #     %timeit shaking_tagger.spots2shaking_tag4batch(batch_matrix_spots) #0.3s\n",
    "    \n",
    "#     for batch_idx in range(len(batch_data)):\n",
    "#         gold_sample = batch_data[batch_idx]\n",
    "#         shaking_tag = batch_shaking_tag[batch_idx]\n",
    "#         # decode\n",
    "#         text = batch_data[batch_idx][\"text\"]\n",
    "#         tok2char_span = get_tok2char_span_map(text)\n",
    "#         rel_list = handshaking_tagger.decode_rel(text, shaking_tag, tok2char_span)\n",
    "#         pred_sample = {\n",
    "#             \"text\": text,\n",
    "#             \"id\": gold_sample[\"id\"],\n",
    "#             \"relation_list\": rel_list,\n",
    "#         }\n",
    "        \n",
    "        \n",
    "#         if not sample_equal_to(pred_sample, gold_sample) or not sample_equal_to(gold_sample, pred_sample):\n",
    "#             set_trace()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if config[\"encoder\"] == \"BERT\":\n",
    "    tokenizer = BertTokenizerFast.from_pretrained(config[\"bert_path\"], add_special_tokens = False, do_lower_case = False)\n",
    "    data_maker = DataMaker4Bert(tokenizer, handshaking_tagger)\n",
    "    \n",
    "elif config[\"encoder\"] in {\"BiLSTM\", }:\n",
    "    token2idx_path = os.path.join(data_home, experiment_name, config[\"token2idx\"])\n",
    "    token2idx = json.load(open(token2idx_path, \"r\", encoding = \"utf-8\"))\n",
    "    idx2token = {idx:tok for tok, idx in token2idx.items()}\n",
    "    def text2indices(text, max_seq_len):\n",
    "        input_ids = []\n",
    "        tokens = text.split(\" \")\n",
    "        for tok in tokens:\n",
    "            if tok not in token2idx:\n",
    "                input_ids.append(token2idx['<UNK>'])\n",
    "            else:\n",
    "                input_ids.append(token2idx[tok])\n",
    "        if len(input_ids) < max_seq_len:\n",
    "            input_ids.extend([token2idx['<PAD>']] * (max_seq_len - len(input_ids)))\n",
    "        input_ids = torch.tensor(input_ids[:max_seq_len])\n",
    "        return input_ids\n",
    "    data_maker = DataMaker4BiLSTM(text2indices, get_tok2char_span_map, handshaking_tagger)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyDataset(Dataset):\n",
    "    def __init__(self, data):\n",
    "        self.data = data\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        return self.data[index]\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "indexed_train_data = data_maker.get_indexed_data(train_data, max_seq_len)\n",
    "indexed_valid_data = data_maker.get_indexed_data(valid_data, max_seq_len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataloader = DataLoader(MyDataset(indexed_train_data), \n",
    "                                  batch_size = hyper_parameters[\"batch_size\"], \n",
    "                                  shuffle = True, \n",
    "                                  num_workers = 6,\n",
    "                                  drop_last = False,\n",
    "                                  collate_fn = data_maker.generate_batch,\n",
    "                                 )\n",
    "valid_dataloader = DataLoader(MyDataset(indexed_valid_data), \n",
    "                          batch_size = hyper_parameters[\"batch_size\"], \n",
    "                          shuffle = True, \n",
    "                          num_workers = 6,\n",
    "                          drop_last = False,\n",
    "                          collate_fn = data_maker.generate_batch,\n",
    "                         )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # have a look at dataloader\n",
    "# train_data_iter = iter(train_dataloader)\n",
    "# batch_data = next(train_data_iter)\n",
    "# text_id_list, text_list, batch_input_ids, \\\n",
    "# batch_attention_mask, batch_token_type_ids, \\\n",
    "# tok2char_span_list, batch_shaking_tag = batch_data\n",
    "\n",
    "# print(text_list[0])\n",
    "# print()\n",
    "# print(tokenizer.decode(batch_input_ids[0].tolist()))\n",
    "# print(batch_input_ids.size())\n",
    "# print(batch_attention_mask.size())\n",
    "# print(batch_token_type_ids.size())\n",
    "# print(len(tok2char_span_list))\n",
    "# print(batch_shaking_tag.size())\n",
    "\n",
    "# # decode \n",
    "# idx = 2\n",
    "# print(text_list[idx])\n",
    "# shaking_tag = batch_shaking_tag[idx]\n",
    "# text = text_list[idx]\n",
    "# tok2char_span = tok2char_span_list[idx]\n",
    "# handshaking_tagger.decode_rel(text, shaking_tag, tok2char_span)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if config[\"encoder\"] == \"BERT\":\n",
    "    encoder = BertModel.from_pretrained(config[\"bert_path\"])\n",
    "    hidden_size = encoder.config.hidden_size\n",
    "    \n",
    "    rel_extractor = TPLinkerPlusBert(encoder, \n",
    "                                     tag_size,\n",
    "                                     hyper_parameters[\"shaking_type\"],\n",
    "                                     hyper_parameters[\"inner_enc_type\"],\n",
    "                                     hyper_parameters[\"tok_pair_sample_rate\"]\n",
    "                                    )\n",
    "    \n",
    "elif config[\"encoder\"] in {\"BiLSTM\", }:\n",
    "    glove = Glove()\n",
    "    glove = glove.load(config[\"pretrained_word_embedding_path\"])\n",
    "    \n",
    "    # prepare embedding matrix\n",
    "    word_embedding_init_matrix = np.random.normal(-1, 1, size=(len(token2idx), hyper_parameters[\"word_embedding_dim\"]))\n",
    "    count_in = 0\n",
    "\n",
    "    # 在预训练词向量中的用该预训练向量\n",
    "    # 不在预训练集里的用随机向量\n",
    "    for ind, tok in tqdm(idx2token.items(), desc=\"Embedding matrix initializing...\"):\n",
    "        if tok in glove.dictionary:\n",
    "            count_in += 1\n",
    "            word_embedding_init_matrix[ind] = glove.word_vectors[glove.dictionary[tok]]\n",
    "\n",
    "    print(\"{:.4f} tokens are in the pretrain word embedding matrix\".format(count_in / len(idx2token))) # 命中预训练词向量的比例\n",
    "    word_embedding_init_matrix = torch.FloatTensor(word_embedding_init_matrix)\n",
    "    \n",
    "    rel_extractor = TPLinkerPlusBiLSTM(word_embedding_init_matrix, \n",
    "                                       hyper_parameters[\"emb_dropout\"], \n",
    "                                       hyper_parameters[\"enc_hidden_size\"], \n",
    "                                       hyper_parameters[\"dec_hidden_size\"],\n",
    "                                       hyper_parameters[\"rnn_dropout\"],\n",
    "                                       tag_size, \n",
    "                                       hyper_parameters[\"shaking_type\"],\n",
    "                                       hyper_parameters[\"inner_enc_type\"],\n",
    "                                       hyper_parameters[\"tok_pair_sample_rate\"],\n",
    "                                      )\n",
    "\n",
    "rel_extractor = rel_extractor.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # test outputs\n",
    "# rel_extractor.train()\n",
    "# with torch.no_grad():\n",
    "#     outputs, sampled_tok_pair_indices = rel_extractor(batch_input_ids.to(device), \n",
    "#                                    batch_attention_mask.to(device), \n",
    "#                                    batch_token_type_ids.to(device),\n",
    "#                                 )\n",
    "# print(outputs.size())\n",
    "# if rel_extractor.training:\n",
    "#     print(sampled_tok_pair_indices.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = MetricsCalculator(handshaking_tagger)\n",
    "loss_func = lambda y_pred, y_true: metrics.loss_func(y_pred, y_true, ghm = hyper_parameters[\"ghm\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train step\n",
    "def train_step(batch_train_data, optimizer):\n",
    "    if config[\"encoder\"] == \"BERT\":\n",
    "        sample_list, batch_input_ids, \\\n",
    "        batch_attention_mask, batch_token_type_ids, \\\n",
    "        tok2char_span_list, batch_shaking_tag = batch_train_data\n",
    "        \n",
    "        batch_input_ids, \\\n",
    "        batch_attention_mask, \\\n",
    "        batch_token_type_ids, \\\n",
    "        batch_shaking_tag = (batch_input_ids.to(device), \n",
    "                              batch_attention_mask.to(device), \n",
    "                              batch_token_type_ids.to(device), \n",
    "                              batch_shaking_tag.to(device)\n",
    "                             )\n",
    "        \n",
    "    elif config[\"encoder\"] in {\"BiLSTM\", }:\n",
    "        sample_list, batch_input_ids, \\\n",
    "        tok2char_span_list, batch_shaking_tag = batch_train_data\n",
    "        \n",
    "        batch_input_ids, \\\n",
    "        batch_shaking_tag = (batch_input_ids.to(device), \n",
    "                              batch_shaking_tag.to(device)\n",
    "                            )\n",
    "    \n",
    "\n",
    "    # zero the parameter gradients\n",
    "    optimizer.zero_grad()\n",
    "    if config[\"encoder\"] == \"BERT\":\n",
    "        pred_small_shaking_outputs, sampled_tok_pair_indices = rel_extractor(batch_input_ids, \n",
    "                                                                            batch_attention_mask, \n",
    "                                                                            batch_token_type_ids\n",
    "                                                                             )\n",
    "    elif config[\"encoder\"] in {\"BiLSTM\", }:\n",
    "        pred_small_shaking_outputs, sampled_tok_pair_indices = rel_extractor(batch_input_ids)\n",
    "\n",
    "    # sampled_tok_pair_indices: (batch_size, ~segment_len)\n",
    "    # batch_small_shaking_tag: (batch_size, ~segment_len, tag_size)\n",
    "    batch_small_shaking_tag = batch_shaking_tag.gather(1, sampled_tok_pair_indices[:, :, None].repeat(1, 1, tag_size))\n",
    "    \n",
    "    loss = loss_func(pred_small_shaking_outputs, batch_small_shaking_tag)\n",
    "#     set_trace()\n",
    "#     t1 = time.time()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "#     print(\"bp: {}\".format(time.time() - t1))\n",
    "    pred_small_shaking_tag = (pred_small_shaking_outputs > 0.).long()\n",
    "    sample_acc = metrics.get_sample_accuracy(pred_small_shaking_tag,\n",
    "                                             batch_small_shaking_tag)\n",
    "    \n",
    "    return loss.item(), sample_acc.item()\n",
    "\n",
    "# valid step\n",
    "def valid_step(batch_valid_data):\n",
    "    if config[\"encoder\"] == \"BERT\":\n",
    "        sample_list, batch_input_ids, \\\n",
    "        batch_attention_mask, batch_token_type_ids, \\\n",
    "        tok2char_span_list, batch_shaking_tag = batch_valid_data\n",
    "        \n",
    "        batch_input_ids, \\\n",
    "        batch_attention_mask, \\\n",
    "        batch_token_type_ids, \\\n",
    "        batch_shaking_tag = (batch_input_ids.to(device), \n",
    "                              batch_attention_mask.to(device), \n",
    "                              batch_token_type_ids.to(device), \n",
    "                              batch_shaking_tag.to(device)\n",
    "                             )\n",
    "        \n",
    "    elif config[\"encoder\"] in {\"BiLSTM\", }:\n",
    "        sample_list, batch_input_ids, \\\n",
    "        tok2char_span_list, batch_shaking_tag = batch_valid_data\n",
    "        \n",
    "        batch_input_ids, \\\n",
    "        batch_shaking_tag = (batch_input_ids.to(device), \n",
    "                              batch_shaking_tag.to(device)\n",
    "                             )\n",
    "    with torch.no_grad():\n",
    "        if config[\"encoder\"] == \"BERT\":\n",
    "            pred_shaking_outputs, _ = rel_extractor(batch_input_ids, \n",
    "                                      batch_attention_mask, \n",
    "                                      batch_token_type_ids, \n",
    "                                     )\n",
    "        elif config[\"encoder\"] in {\"BiLSTM\", }:\n",
    "            pred_shaking_outputs, _ = rel_extractor(batch_input_ids)\n",
    "            \n",
    "    pred_shaking_tag = (pred_shaking_outputs > 0.).long()\n",
    "    sample_acc = metrics.get_sample_accuracy(pred_shaking_tag,\n",
    "                                             batch_shaking_tag)\n",
    "    \n",
    "    cpg_dict = metrics.get_cpg(sample_list, \n",
    "                                          tok2char_span_list, \n",
    "                                          pred_shaking_tag,\n",
    "                                          hyper_parameters[\"match_pattern\"])\n",
    "    return sample_acc.item(), cpg_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_f1 = 0.\n",
    "def train_n_valid(train_dataloader, dev_dataloader, optimizer, scheduler, num_epoch):  \n",
    "    def train(dataloader, ep):\n",
    "        # train\n",
    "        rel_extractor.train()\n",
    "        \n",
    "        t_ep = time.time()\n",
    "        total_loss, total_sample_acc = 0., 0.\n",
    "        for batch_ind, batch_train_data in enumerate(dataloader):\n",
    "            t_batch = time.time()\n",
    "            \n",
    "            loss, sample_acc = train_step(batch_train_data, optimizer)\n",
    "            \n",
    "            total_loss += loss\n",
    "            total_sample_acc += sample_acc\n",
    "            \n",
    "            avg_loss = total_loss / (batch_ind + 1)\n",
    "            \n",
    "            # scheduler\n",
    "            if hyper_parameters[\"scheduler\"] == \"ReduceLROnPlateau\":\n",
    "                scheduler.step(avg_loss)\n",
    "            else:\n",
    "                scheduler.step()\n",
    "                \n",
    "            avg_sample_acc = total_sample_acc / (batch_ind + 1)\n",
    "            \n",
    "            batch_print_format = \"\\rproject: {}, run_name: {}, Epoch: {}/{}, batch: {}/{}, train_loss: {}, \" + \\\n",
    "                                \"t_sample_acc: {},\" + \\\n",
    "                                 \"lr: {}, batch_time: {}, total_time: {} -------------\"\n",
    "                    \n",
    "            print(batch_print_format.format(experiment_name, config[\"run_name\"], \n",
    "                                            ep + 1, num_epoch, \n",
    "                                            batch_ind + 1, len(dataloader), \n",
    "                                            avg_loss, \n",
    "                                            avg_sample_acc,\n",
    "                                            optimizer.param_groups[0]['lr'],\n",
    "                                            time.time() - t_batch,\n",
    "                                            time.time() - t_ep,\n",
    "                                           ), end=\"\")\n",
    "            \n",
    "            if config[\"logger\"] == \"wandb\" and batch_ind % hyper_parameters[\"log_interval\"] == 0:\n",
    "                logger.log({\n",
    "                    \"train_loss\": avg_loss,\n",
    "                    \"train_small_shaking_seq_acc\": avg_sample_acc,\n",
    "                    \"learning_rate\": optimizer.param_groups[0]['lr'],\n",
    "                    \"time\": time.time() - t_ep,\n",
    "                })\n",
    "                \n",
    "        if config[\"logger\"] != \"wandb\": # only log once for training if logger is not wandb\n",
    "                logger.log({\n",
    "                    \"train_loss\": avg_loss,\n",
    "                    \"train_small_shaking_seq_acc\": avg_sample_acc,\n",
    "                    \"learning_rate\": optimizer.param_groups[0]['lr'],\n",
    "                    \"time\": time.time() - t_ep,\n",
    "                }) \n",
    "            \n",
    "        \n",
    "    def valid(dataloader, ep):\n",
    "        # valid\n",
    "        rel_extractor.eval()\n",
    "        \n",
    "        t_ep = time.time()\n",
    "        total_sample_acc = 0.\n",
    "#         total_rel_correct_num, total_rel_pred_num, total_rel_gold_num = 0, 0, 0\n",
    "#         total_ent_correct_num, total_ent_pred_num, total_ent_gold_num = 0, 0, 0\n",
    "        total_cpg_dict = {}\n",
    "        for batch_ind, batch_valid_data in enumerate(tqdm(dataloader, desc = \"Validating\")):\n",
    "            sample_acc, cpg_dict = valid_step(batch_valid_data)\n",
    "            total_sample_acc += sample_acc\n",
    "            \n",
    "            # init total_cpg_dict\n",
    "            for k in cpg_dict.keys():\n",
    "                if k not in total_cpg_dict:\n",
    "                    total_cpg_dict[k] = [0, 0, 0]\n",
    "           \n",
    "            for k, cpg in cpg_dict.items():\n",
    "                for idx, n in enumerate(cpg):\n",
    "                    total_cpg_dict[k][idx] += cpg[idx]\n",
    "\n",
    "        avg_sample_acc = total_sample_acc / len(dataloader)\n",
    "        \n",
    "        if \"rel_cpg\" in total_cpg_dict:\n",
    "            rel_prf = metrics.get_prf_scores(total_cpg_dict[\"rel_cpg\"][0], total_cpg_dict[\"rel_cpg\"][1], total_cpg_dict[\"rel_cpg\"][2])\n",
    "            ent_prf = metrics.get_prf_scores(total_cpg_dict[\"ent_cpg\"][0], total_cpg_dict[\"ent_cpg\"][1], total_cpg_dict[\"ent_cpg\"][2])\n",
    "            final_score = rel_prf[2]\n",
    "            log_dict = {\n",
    "                            \"val_shaking_tag_acc\": avg_sample_acc,\n",
    "                            \"val_rel_prec\": rel_prf[0],\n",
    "                            \"val_rel_recall\": rel_prf[1],\n",
    "                            \"val_rel_f1\": rel_prf[2],\n",
    "                            \"val_ent_prec\": ent_prf[0],\n",
    "                            \"val_ent_recall\": ent_prf[1],\n",
    "                            \"val_ent_f1\": ent_prf[2],\n",
    "                            \"time\": time.time() - t_ep,\n",
    "                        }\n",
    "        elif \"trigger_iden_cpg\" in total_cpg_dict:\n",
    "            trigger_iden_prf = metrics.get_prf_scores(total_cpg_dict[\"trigger_iden_cpg\"][0], \n",
    "                                                      total_cpg_dict[\"trigger_iden_cpg\"][1], \n",
    "                                                      total_cpg_dict[\"trigger_iden_cpg\"][2])\n",
    "            trigger_class_prf = metrics.get_prf_scores(total_cpg_dict[\"trigger_class_cpg\"][0], \n",
    "                                                      total_cpg_dict[\"trigger_class_cpg\"][1], \n",
    "                                                      total_cpg_dict[\"trigger_class_cpg\"][2])\n",
    "            arg_iden_prf = metrics.get_prf_scores(total_cpg_dict[\"arg_iden_cpg\"][0], total_cpg_dict[\"arg_iden_cpg\"][1], total_cpg_dict[\"arg_iden_cpg\"][2])\n",
    "            arg_class_prf = metrics.get_prf_scores(total_cpg_dict[\"arg_class_cpg\"][0], total_cpg_dict[\"arg_class_cpg\"][1], total_cpg_dict[\"arg_class_cpg\"][2])\n",
    "            final_score = arg_class_prf[2]\n",
    "            log_dict = {\n",
    "                            \"val_shaking_tag_acc\": avg_sample_acc,\n",
    "                            \"val_trigger_iden_prec\": trigger_iden_prf[0],\n",
    "                            \"val_trigger_iden_recall\": trigger_iden_prf[1],\n",
    "                            \"val_trigger_iden_f1\": trigger_iden_prf[2],\n",
    "                            \"val_trigger_class_prec\": trigger_class_prf[0],\n",
    "                            \"val_trigger_class_recall\": trigger_class_prf[1],\n",
    "                            \"val_trigger_class_f1\": trigger_class_prf[2],\n",
    "                            \"val_arg_iden_prec\": arg_iden_prf[0],\n",
    "                            \"val_arg_iden_recall\": arg_iden_prf[1],\n",
    "                            \"val_arg_iden_f1\": arg_iden_prf[2],\n",
    "                            \"val_arg_class_prec\": arg_class_prf[0],\n",
    "                            \"val_arg_class_recall\": arg_class_prf[1],\n",
    "                            \"val_arg_class_f1\": arg_class_prf[2],\n",
    "                            \"time\": time.time() - t_ep,\n",
    "                        }\n",
    "            \n",
    "        logger.log(log_dict)\n",
    "        pprint(log_dict)\n",
    "        \n",
    "        return final_score\n",
    "        \n",
    "    for ep in range(num_epoch):\n",
    "        train(train_dataloader, ep)   \n",
    "        valid_f1 = valid(valid_dataloader, ep)\n",
    "        \n",
    "        global max_f1\n",
    "        if valid_f1 >= max_f1: \n",
    "            max_f1 = valid_f1\n",
    "            if valid_f1 > config[\"f1_2_save\"]: # save the best model\n",
    "                modle_state_num = len(glob.glob(model_state_dict_dir + \"/model_state_dict_*.pt\"))\n",
    "                torch.save(rel_extractor.state_dict(), os.path.join(model_state_dict_dir, \"model_state_dict_{}.pt\".format(modle_state_num)))\n",
    "#                 scheduler_state_num = len(glob.glob(schedule_state_dict_dir + \"/scheduler_state_dict_*.pt\"))\n",
    "#                 torch.save(scheduler.state_dict(), os.path.join(schedule_state_dict_dir, \"scheduler_state_dict_{}.pt\".format(scheduler_state_num))) \n",
    "        print(\"Current avf_f1: {}, Best f1: {}\".format(valid_f1, max_f1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# optimizer\n",
    "init_learning_rate = float(hyper_parameters[\"lr\"])\n",
    "optimizer = torch.optim.Adam(rel_extractor.parameters(), lr = init_learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if hyper_parameters[\"scheduler\"] == \"CAWR\":\n",
    "    T_mult = hyper_parameters[\"T_mult\"]\n",
    "    rewarm_epoch_num = hyper_parameters[\"rewarm_epoch_num\"]\n",
    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, len(train_dataloader) * rewarm_epoch_num, T_mult)\n",
    "    \n",
    "elif hyper_parameters[\"scheduler\"] == \"Step\":\n",
    "    decay_rate = hyper_parameters[\"decay_rate\"]\n",
    "    decay_steps = hyper_parameters[\"decay_steps\"]\n",
    "    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size = decay_steps, gamma = decay_rate)\n",
    "\n",
    "elif hyper_parameters[\"scheduler\"] == \"ReduceLROnPlateau\":\n",
    "    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, \"min\", verbose = True, patience = 6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if not config[\"fr_scratch\"]:\n",
    "    model_state_path = config[\"model_state_dict_path\"]\n",
    "    rel_extractor.load_state_dict(torch.load(model_state_path))\n",
    "    print(\"------------model state {} loaded ----------------\".format(model_state_path.split(\"/\")[-1]))\n",
    "\n",
    "\n",
    "train_n_valid(train_dataloader, valid_dataloader, optimizer, scheduler, hyper_parameters[\"epochs\"])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
